"""标定位移平台坐标系到图形坐标系的关系"""
from pyabb import XYZPlatform
from pyati import ATISensor
from gelslimpi import GelSlimInterface, InputSource
from gelslimpi.visualize import qtcv, QImageViewWidget, QTimer, Qt

import cv2
import sys
from PyQt5 import QtGui, QtWidgets, QtCore
from scipy.interpolate import LinearNDInterpolator
from scipy.interpolate import RegularGridInterpolator as RGI

from datetime import datetime

from time import sleep
import numpy as np
from pathlib import Path
PROJ_DIR = Path(__file__).resolve().parent
CONF_DIR = PROJ_DIR / "configs"


class XYZ_Cali(QtWidgets.QWidget):
    """移动到任意像素坐标
    采集图片
    采集传感器数据
    下压"""
    def __init__(self, z_cont):
        super().__init__()
        self.xyz = XYZPlatform("/dev/ttyUSB0", [-57, -119, -120])
        self.ati = ATISensor()
        self.ati.tare()
        self.source = InputSource("/raspicam_node32/image/compressed", "gelslim")
        self.gelslim = GelSlimInterface(configs=CONF_DIR/"configs3.yaml", input_source=self.source)
        self.init_interp()

        self.x_pix, self.y_pix = -1, -1
        self.point_pix = []
        self.point_xyz = []
        self.pts_gen = self.points_generator()

        self.init_ui()
        self.timer = QtCore.QTimer(self)
        self.timer.start(100)
        self.timer.timeout.connect(self.on_timeout)

        self.z_cont = z_cont
        self.cont_th1 = -0.015
        self.cont_th2 = -0.035


    def init_ui(self):
        self.resize(640, 480)
        hbox = QtWidgets.QHBoxLayout(self)
        self.image_view = QImageViewWidget(self)
        hbox.addWidget(self.image_view)

    def on_timeout(self):
        img = self.get_frame()
        cv2.circle(img, (int(self.x_pix), int(self.y_pix)), 5, (0, 0, 255), 2)
        cv2.circle(img, (int(self.x_pix), int(self.y_pix)), 25, (0, 225, 225), 1)
        self.image_view.setData(img)

    def keyPressEvent(self, ev) -> None:
        """按键控制
        enter: 下一个标定点
        space: 随机一个测试点
        left/right/up/down: 调整图中框选位置
        escape: 打印结果， 退出
        0: 移动到(320, 240), 向下移动， 并且不断打印力的数值
        """
        key = ev.key()
        if key in [Qt.Key.Key_Return, Qt.Key.Key_Enter]:
            # 保存当前位置
            if self.x_pix != -1:
                x, y, z = self.get_xyz()
                self.point_xyz.append([x, y])
                self.point_pix.append([self.x_pix, self.y_pix])
                print(f"save {self.x_pix}, {self.y_pix}")
            # 移动到下一个位置
            self.x_pix, self.y_pix = next(self.pts_gen)
            self.move(self.x_pix, self.y_pix)
        elif key == Qt.Key.Key_Space:
            if len(self.point_pix) > 10:
                self.init_interp(self.point_pix, self.point_xyz)
            self.x_pix, self.y_pix = x = np.random.randint(20, 620), np.random.randint(20, 460)
            self.move(self.x_pix, self.y_pix)
        elif key == Qt.Key.Key_Left:
            self.x_pix -= 1
        elif key == Qt.Key.Key_Right:
            self.x_pix += 1
        elif key == Qt.Key.Key_Up:
            self.y_pix -= 1
        elif key == Qt.Key.Key_Down:
            self.y_pix += 1
        elif key == Qt.Key.Key_Escape:
            print(self.point_pix)
            print(self.point_xyz)
            exit()
        elif key == Qt.Key.Key_0:
            self.move_down()

    def get_frame(self):
        return self.gelslim.get_frame()[0]

    def get_ati(self):
        return self.ati.data

    def get_xyz(self):
        return self.xyz.x, self.xyz.y, self.xyz.z

    def init_interp(self, points_pix=None, points_xyz=None):
        """将像素坐标映射为实际坐标"""
        if points_pix is None:
            points_pix = np.array(
        [[45, 39], [315, 41], [591, 38], [580, 230], [577, 421], [313, 431], [43, 421], [41, 228], [314, 227], [549, 394]]
                                  )
        else:
            points_pix = np.array(points_pix)
        if points_xyz is None:
            points_xyz = np.array(
[[9.625, -63.286], [1.645, -63.56], [-6.481, -63.563], [-6.401, -57.836], [-6.51, -52.211], [1.53, -52.025], [9.417, -51.936], [9.67, -57.657], [1.631, -57.898], [-5.61, -53.083]]
)
        else:
            points_xyz = np.array(points_xyz)
        # 用非结构化插值获得网格数据
        interp = LinearNDInterpolator(points=points_pix, values=points_xyz)
        ymin = np.max(points_pix[0:3,1])
        xmax = np.min(points_pix[2:5,0])
        ymax = np.min(points_pix[4:7,1])
        xmin = np.max([points_pix[6,0], points_pix[7,0], points_pix[0,0]])
        x_span = np.linspace(xmin, xmax, 9)
        y_span = np.linspace(ymin, ymax, 9)
        X, Y = np.meshgrid(x_span, y_span, indexing="ij")
        grids = interp(X, Y)
        self.interp = RGI((x_span, y_span), grids, bounds_error=False, fill_value=None)

    def move(self, pix_x, pix_y):
        # 移动到接触高度上方， 力传感器清零
        self.xyz.move_abs(z = self.z_cont+2, vel=30)  # z 提升到 -26
        self.ati.tare()
        # xy 移动
        x, y = self.interp((pix_x, pix_y))
        print(f"move to [{pix_x:.2f}, {pix_y:.2f}]pix, [{x:.2f}, {y:.2f}]mm")
        self.xyz.move_abs(x=x, y=y, vel=30)
        # 向下移动
        self.xyz.move_rel(z=-1.5)  # for speed
        self.down_util_contact()
        # 直接向下移动一段距离， 使得接触明显
        self.xyz.move_rel(z=-0.8, vel=30)


    def down_util_contact(self):
        cont_cnt = 0  # 检测到接触的次数
        step = [0.2, 0.01]
        while cont_cnt < len(step):
            # move down slowly
            self.xyz.move_rel(z=-step[cont_cnt], vel=40)
            sleep(0.05)
            img = self.get_frame()
            qtcv.imshow("contact", img)
            key = qtcv.waitKey(5)
            if key == qtcv.Key.Key_Escape:
                exit()

            # 检测法向力
            fz1 = self.get_ati()[2]
            if fz1 <= self.cont_th1:
                # 再向下 0.1 确认是否接触
                self.xyz.move_rel(z=-0.1, vel=40)
                fz2 = self.get_ati()[2]
                is_contact = fz2 <= self.cont_th2
                self.xyz.move_rel(z=0.1, vel=40)
                if is_contact and cont_cnt == len(step)-1:  # 最后一次收敛
                    self.z_cont = self.xyz.z
                    print(f"contact detected. z_cont: {self.z_cont}mm | fz1: {fz1} | fz2: {fz2}")
                    break
                elif is_contact:  # 不是最后一次接触
                    self.xyz.move_rel(z=step[cont_cnt], vel=40)
                    cont_cnt += 1

    def is_contacted(self, threshold):
        fz = self.get_ati()[2]
        # print(fz)
        return fz <= threshold

    def points_generator(self):
        points = [[50, 50], [320,50], [590, 50],
               [590, 240], [590, 430],
               [320, 430], [50, 430],
               [50, 240],
               [320, 240], [560, 400]]
        i = 0
        while True:
            if i < len(points):
                yield points[i]
                i += 1
            else:
                yield [np.random.randint(20, 620), np.random.randint(20, 460)]


if __name__ == '__main__':
    import sys
    app = QtWidgets.QApplication(sys.argv)
    win = XYZ_Cali(z_cont=-235)
    win.show()
    sys.exit(app.exec_())

